package com.atguigu.guli.ai.strategy;

import com.atguigu.guli.ai.constant.SystemConstant;

import com.atguigu.guli.ai.vo.QueryVo;
import com.atguigu.system.domain.ChatKnowledge;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

@AiBean(SystemConstant.MODEL_TYPE_OLLAMA)
public class OllamaAiOperator implements AiOperator{
    @Autowired
    private QdrantVectorStore ollamaVectorStore;
    @Autowired
    private OllamaChatModel ollamaChatModel;

    @Override
    public void addDosc(ChatKnowledge chatKnowledge) {
        ollamaVectorStore.add(List.of(new Document(chatKnowledge.getContent(),
                Map.of("projectId",chatKnowledge.getProjectId().toString(),"knowledgeId",chatKnowledge.getKnowledgeId().toString()))));
    }

    @Override
    public List<Document> similaritySearch(QueryVo queryVo) {
        List<Document> docs = ollamaVectorStore.similaritySearch(
                SearchRequest.query(queryVo.getMsg())
                        .withFilterExpression(new FilterExpressionBuilder()
                                .eq("projectId", queryVo.getProjectId().toString())
                                .build())
                        .withTopK(SystemConstant.TOP_K));
        return docs;
    }

    @Override
    public Flux<String> chat_stream(List<Message> msgs) {
        Flux<String> result = this.ollamaChatModel.stream(msgs.toArray(new Message[]{}));
        return result;
    }
}
